import { streamSse } from "../../../fetch/stream.js"
import {
	ChatCompletion,
	ChatCompletionChunk,
	ChatCompletionCreateParamsNonStreaming,
	ChatCompletionCreateParamsStreaming,
	ChatCompletionMessageParam,
	Model,
} from "openai/resources/index"
import { InceptionConfig } from "../types.js"
import { chatChunk } from "../util.js"
import { OpenAIApi } from "./OpenAI.js"
import { FimCreateParamsStreaming } from "./base.js"

// export type InceptionChatCompletionCreateParamsStreaming =
//   ChatCompletionCreateParamsStreaming & {
//     nextEdit?: boolean;
//   };

// export type InceptionChatCompletionCreateParamsNonStreaming =
//   ChatCompletionCreateParamsNonStreaming & {
//     nextEdit?: boolean;
//   };

const UNIQUE_TOKEN = "<|!@#IS_NEXT_EDIT!@#|>"
const APPLY_UNIQUE_TOKEN = "<|!@#IS_APPLY!@#|>"
const INCEPTION_API_BASE = "https://api.inceptionlabs.ai/v1/"
export class InceptionApi extends OpenAIApi {
	constructor(config: InceptionConfig) {
		super({
			...config,
			provider: "openai",
			apiBase: config.apiBase ?? INCEPTION_API_BASE,
		})
	}

	// Add custom edit completions method.
	async *editCompletionStream(
		body: ChatCompletionCreateParamsStreaming,
		signal: AbortSignal,
	): AsyncGenerator<ChatCompletionChunk, any, unknown> {
		yield* this.streamCustomEndpoint("edit/completions", body, signal)
	}

	// Add custom edit completions method (non-streaming).
	async editCompletionNonStream(
		body: ChatCompletionCreateParamsNonStreaming,
		signal: AbortSignal,
	): Promise<ChatCompletion> {
		return this.nonStreamCustomEndpoint("edit/completions", body, signal)
	}

	// Override the regular chat stream method to route to edit endpoint for next edit requests.
	override async *chatCompletionStream(
		body: ChatCompletionCreateParamsStreaming,
		signal: AbortSignal,
	): AsyncGenerator<ChatCompletionChunk, any, unknown> {
		if (this.isNextEdit(body.messages)) {
			body.messages = this.removeToken(body.messages, UNIQUE_TOKEN)
			yield* this.editCompletionStream(body, signal)
		} else if (this.isApply(body.messages)) {
			body.messages = this.removeToken(body.messages, APPLY_UNIQUE_TOKEN)
			yield* this.applyCompletionStream(body, signal)
		} else {
			yield* super.chatCompletionStream(body, signal)
		}
	}

	// Override the regular chat non stream method to route to edit endpoint for next edit requests.
	override async chatCompletionNonStream(
		body: ChatCompletionCreateParamsNonStreaming,
		signal: AbortSignal,
	): Promise<ChatCompletion> {
		if (this.isNextEdit(body.messages)) {
			body.messages = this.removeToken(body.messages, UNIQUE_TOKEN)
			return this.editCompletionNonStream(body, signal)
		} else if (this.isApply(body.messages)) {
			body.messages = this.removeToken(body.messages, APPLY_UNIQUE_TOKEN)
			return this.applyCompletionNonStream(body, signal)
		} else {
			return super.chatCompletionNonStream(body, signal)
		}
	}

	async *applyCompletionStream(
		body: ChatCompletionCreateParamsStreaming,
		signal: AbortSignal,
	): AsyncGenerator<ChatCompletionChunk, any, unknown> {
		yield* this.streamCustomEndpoint("apply/completions", body, signal)
	}

	async applyCompletionNonStream(
		body: ChatCompletionCreateParamsNonStreaming,
		signal: AbortSignal,
	): Promise<ChatCompletion> {
		return this.nonStreamCustomEndpoint("apply/completions", body, signal)
	}

	override async *fimStream(
		body: FimCreateParamsStreaming,
		signal: AbortSignal,
	): AsyncGenerator<ChatCompletionChunk, any, unknown> {
		const endpoint = new URL("completions", this.apiBase)
		const resp = await fetch(endpoint, {
			method: "POST",
			body: JSON.stringify({
				model: body.model,
				prompt: body.prompt,
				suffix: body.suffix.trim() === "" ? "<|endoftext|>" : body.suffix,
				max_tokens: body.max_tokens ?? 150, // Only want this for /fim, not chat
				temperature: body.temperature,
				top_p: body.top_p,
				frequency_penalty: body.frequency_penalty,
				presence_penalty: body.presence_penalty,
				stop: [...(body.stop ?? []), "\n\n", "\n \n"],
				stream: true,
			}),
			headers: {
				"Content-Type": "application/json",
				Accept: "application/json",
				Authorization: `Bearer ${this.config.apiKey}`,
			},
			signal,
		})
		for await (const chunk of streamSse(resp as any)) {
			if (!chunk.choices[0]) {
				continue
			}
			yield chatChunk({
				content: chunk.choices[0].text,
				finish_reason: null,
				model: body.model,
			})
		}
	}

	override list(): Promise<Model[]> {
		throw new Error("Method not implemented.")
	}

	// Check if any message contains the unique next edit token.
	private isNextEdit(messages: ChatCompletionMessageParam[]): boolean {
		return messages.some((message) => typeof message.content === "string" && message.content.endsWith(UNIQUE_TOKEN))
	}

	private isApply(messages: ChatCompletionMessageParam[]): boolean {
		return messages.some(
			(message) => typeof message.content === "string" && message.content.endsWith(APPLY_UNIQUE_TOKEN),
		)
	}

	// Remove the unique token from messages.
	private removeToken(messages: ChatCompletionMessageParam[], token: string): ChatCompletionMessageParam[] {
		const lastMessage = messages[messages.length - 1]

		if (typeof lastMessage?.content === "string" && lastMessage.content.endsWith(token)) {
			const cleanedMessages = [...messages]
			cleanedMessages[cleanedMessages.length - 1] = {
				...lastMessage,
				content: lastMessage.content.slice(0, -token.length),
			}
			return cleanedMessages
		}

		return messages
	}

	private async *streamCustomEndpoint(
		path: string,
		body: ChatCompletionCreateParamsStreaming,
		signal: AbortSignal,
	): AsyncGenerator<ChatCompletionChunk, any, unknown> {
		const endpoint = new URL(path, this.apiBase)
		const resp = await fetch(endpoint, {
			method: "POST",
			body: JSON.stringify({
				model: body.model,
				messages: body.messages,
				max_tokens: body.max_tokens,
				temperature: body.temperature,
				top_p: body.top_p,
				frequency_penalty: body.frequency_penalty,
				presence_penalty: body.presence_penalty,
				stop: body.stop,
				stream: true,
			}),
			headers: {
				"Content-Type": "application/json",
				Accept: "application/json",
				Authorization: `Bearer ${this.config.apiKey}`,
			},
			signal,
		})

		for await (const chunk of streamSse(resp as any)) {
			const deltaContent = chunk.choices?.[0]?.delta?.content
			if (deltaContent) {
				yield chatChunk({
					content: deltaContent,
					finish_reason: chunk.choices[0].finish_reason || null,
					model: body.model,
				})
			}
		}
	}

	private async nonStreamCustomEndpoint(
		path: string,
		body: ChatCompletionCreateParamsNonStreaming,
		signal: AbortSignal,
	): Promise<ChatCompletion> {
		const endpoint = new URL(path, this.apiBase)
		const resp = await fetch(endpoint, {
			method: "POST",
			body: JSON.stringify({
				model: body.model,
				messages: body.messages,
				max_tokens: body.max_tokens,
				temperature: body.temperature,
				top_p: body.top_p,
				frequency_penalty: body.frequency_penalty,
				presence_penalty: body.presence_penalty,
				stop: body.stop,
				stream: false,
			}),
			headers: {
				"Content-Type": "application/json",
				Accept: "application/json",
				Authorization: `Bearer ${this.config.apiKey}`,
			},
			signal,
		})

		const data = await resp.json()
		return data as ChatCompletion
	}
}
